import urequests
import ujson
from setcfg import DAT
import dog_control
import time
import os

import wifi_moudle
wifi_moudle.InitWifi()

#智能体模块
# 新增：记忆文件路径

MAX_MEMORY = 5  # 最大记忆条数

max_tokens = DAT["openai"]["max_tokens"]  # 新增最大token数配置

def load_memory():
    try:
        if MEMORY_FILE in os.listdir():
            with open(MEMORY_FILE, "r") as f:
                return ujson.load(f)
    except:
        pass
    return []

def call_ai_api(prompt):
    if len(prompt) < 2:
        return
    
    # 加载记忆
    memory = load_memory()
    
    api_url = DAT["openai"]["url"]
    if api_url == "" or DAT["openai"]["apikey"] == "":
        print("请配置openai接口和apikey")
        return 
    
    # 构建带有记忆的消息
    messages = [
        {"role": "system", "content": "你的名字叫胖胖狗,是一只聪明的机器狗.#1:当被叫到名字时，会执行名字应答指令.#2:能听从指令执行动作,模糊音、同音词语句优先识别为命令.#3:没有匹配指令的时候,能陪主人聊天"}
    ]
    messages.extend(memory)
    messages.append({"role": "user", "content": prompt})

    headers = {
        "Content-Type": "application/json",
        "Authorization": "Bearer "+DAT["openai"]["apikey"]
    }
    payload = {
        "model": DAT["openai"]["model"],
        "messages": messages,
        "max_tokens": DAT["openai"]["max_tokens"],
        "tools": [
        {
            'type': 'function',
            'function': {
                "name": "hand_cmd",
                "description": "执行本地命令操作,有2个参数,第一个参数为操作类型,第二个参数为运行时间",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "opera": {
                            "type": "string",
                            "description": f"需要执行的操作指令{dog_control.cmdlist}",
                            "enum": dog_control.cmdlist
                        },
                        "rt": {
                            "type": "number",
                            "description": "运行时间,单位为秒,除了停止是0,其他运行时间范围5-20秒"
                        }
                    },
                    "required": ["opera", "rt"]
                },
            }
        },
        {
            'type': 'function',
            'function': {
                "name": "name_replay",
                "description": "名字应答,被叫到胖胖狗这个名字的时候执行这个指令",
                "parameters": {
                    "_comments": "无参数"
                },
            }
        }
        ]        
    }
    
    try:        
        response = urequests.post(api_url, headers=headers, data=ujson.dumps(payload).encode('utf8'))
        #print(response.text)
        result = ujson.loads(response.text)
        if 'choices' in result and result['choices']:
            choice = result['choices'][0]
            total_tokens = result.get('usage', {}).get('total_tokens', 0)
            
            # 修复：给用户消息也添加token字段
            memory.append({
                "role": "user", 
                "content": prompt,
                "tokens": 0  # 用户消息固定0token
            })
            
            if choice['message'].get('content'):
                memory.append({
                    "role": "assistant",
                    "content": choice['message']['content'],
                    "tokens": total_tokens
                })
            
            save_memory(memory)
            if choice['message'].get('tool_calls') and len(choice['message']['tool_calls'])>0:
                for tool_call in choice['message']['tool_calls']:
                    func_call = tool_call["function"]
                    print(f"函数调用: {func_call['name']} 参数: {func_call['arguments']}")
                     #直接在此处处理函数调用
                    func_name = func_call['name']
                    args = ujson.loads(func_call['arguments'])
                    if func_name == 'hand_cmd':
                      hand_cmd(args['opera'],args['rt'])  # 直接调用本地函数
                      time.sleep(args['rt'])
                    elif func_name == 'name_replay':
                      name_replay()
        else:
            speak(choice['message']['content'])
        return None
    except Exception as e:
        print(f"API调用失败: {e}")
        return None

def hand_cmd(opera,rt):
    # 实际执行命令的代码   
    if opera in dog_control.cmdlist:        
       dog_control.hand_cmd(opera,rt) 
    else:
       print(f"无效操作：{opera}")

def name_replay():
    from tts_module_m import play_wav
    play_wav("wangwang.wav")

def speak(text):
    from tts_module_m import to_tts
    to_tts(text)

# 修改文件路径定义（MicroPython通常使用相对路径）
MEMORY_FILE = "memory.json"  # 移除了根目录斜杠

def save_memory(memory):
    try:
        current_total = sum(entry.get('tokens', 0) for entry in memory)
        threshold = max_tokens // 2
        
        if current_total <= threshold:
            with open(MEMORY_FILE, "w") as f:
                ujson.dump(memory, f)
            return
        
        # 修改后的清理逻辑
        needed_deletion = current_total - threshold
        accumulated = 0
        keep_start = len(memory)  # 保留起始位置
        
        i = 0
        while i < len(memory):
            # 处理成对的用户消息和助理消息
            if i+1 < len(memory) and \
               memory[i]['role'] == 'user' and \
               memory[i+1]['role'] == 'assistant':
                
                pair_tokens = memory[i+1].get('tokens', 0)
                accumulated += pair_tokens
                
                if accumulated <= needed_deletion:
                    # 继续删除这对记录
                    i += 2
                else:
                    # 找到保留起始点
                    keep_start = i + 2
                    break
            else:
                i += 1
        
        # 保留从keep_start开始的所有记录
        new_memory = memory[keep_start:]
        
        with open(MEMORY_FILE, "w") as f:
            ujson.dump(new_memory, f)
            
    except Exception as e:
        print(f"保存记忆失败: {e}")

# 其他地方直接调用call_ai_api即可
#call_ai_api("往前走两米")
